/**
 * 
 */
package cs534.crf;

import java.util.Arrays;

import org.apache.commons.math3.random.MersenneTwister;


/**
 * @author jhostetler
 *
 */
public class CrfGibbsSampler<Observation> extends CrfSampler<Observation>
{
	public final int burn_in;
	
	private final MersenneTwister rng_;
	
	private final int[] state_;
	private final int[] variable_order_;
	
	private Observation x_ = null;
	
	/**
	 * @param crf
	 */
	public CrfGibbsSampler( final ConditionalRandomField<Observation> crf, final int burn_in, final long rng_seed ) 
	{
		super( crf );
		this.burn_in = burn_in;
		
		rng_ = new MersenneTwister( rng_seed );
		
		// Uniform random initialization
		state_ = new int[crf_.Nvertices];
		initStateUniform();
		variable_order_ = new int[crf_.Nvertices];
		for( int v = 0; v < crf_.Nvertices; ++v ) {
			variable_order_[v] = v;
		}
	}
	
	private void initStateUniform()
	{
		for( int v = 0; v < crf_.Nvertices; ++v ) {
			final int y = rng_.nextInt( crf_.cardinality( v ) );
			state_[v] = y;
		}
	}
	
	private void burnIn( final int n )
	{
		for( int i = 0; i < n; ++i ) {
			sample();
		}
	}
	
	public void enterObservation( final Observation x )
	{
		x_ = x;
		initStateUniform();
		burnIn( burn_in );
	}
	
	private double[] posterior( final int v, final Observation x )
	{
		final int[] cliques = crf_.cliques( v );
		final double[] potential = new double[crf_.cardinality( v )];
		double Z = 0.0;
		for( int i = 0; i < crf_.cardinality( v ); ++i ) {
			double phi = 0.0;
			for( int c : cliques ) {
				final ConditionalRandomField.Clique<Observation> clique = crf_.cliques()[c];
				final int[] labels = new int[clique.vertices.length];
				for( int cvi = 0; cvi < labels.length; ++cvi ) {
					final int cv = clique.vertices[cvi];
					labels[cvi] = (cv == v ? i : state_[cv]);
				}
//				final double p = clique.logPotential( x, labels );
//				System.out.println( "y = " + labels[0] + "; potential = " + p );
				phi += clique.logPotential( x, labels );
			}
			final double exp_phi = Math.exp( phi );
			Z += exp_phi;
			potential[i] = exp_phi;
		}
		assert( Z > 0 );
		final double log_Z = Math.log( Z );
		for( int i = 0; i < potential.length; ++i ) {
			potential[i] = Math.exp( Math.log( potential[i] ) - log_Z );
		}
		return potential;
	}
	
	private void shuffleInplace( final int[] arr )
	{
		for( int i = arr.length - 1; i >= 1; --i ) {
			final int j = rng_.nextInt( i + 1 );
			final int temp = arr[i];
			arr[i] = arr[j];
			arr[j] = temp;
		}
	}
	
	private int sampleCategorical( final double[] p )
	{
		double test = rng_.nextDouble();
		for( int i = 0; i < p.length - 1; ++i ) {
			test -= p[i];
			if( test < 0 ) {
				return i;
			}
		}
		return p.length - 1;
	}
	
	/**
	 * @see cs534.crf.CrfSampler#sample()
	 */
	@Override
	public int[] sample()
	{
		shuffleInplace( variable_order_ );
//		int idx = 0;
		for( int v : variable_order_ ) {
//			System.out.println( "v[" + idx++ + "] = " + v );
			final double[] post = posterior( v, x_ );
//			System.out.println( "posterior: [" + post[0] + ", " + post[1] + "]" );
			final int v_prime = sampleCategorical( post );
			state_[v] = v_prime;
		}
		return Arrays.copyOf( state_, state_.length );
	}

	@Override
	public CrfSampler<Observation> update( final ConditionalRandomField<Observation> new_crf )
	{
		// I want to pass the same seed, but it's not available from the RNG.
		// Generating a seed from the RNG should at least make it deterministic.
		return new CrfGibbsSampler<Observation>( new_crf, burn_in, rng_.nextInt() );
	}

}
